package ai.onnxruntime.example.objectdetection

import ai.onnxruntime.OnnxJavaType
import ai.onnxruntime.OrtSession
import ai.onnxruntime.OnnxTensor
import ai.onnxruntime.OrtEnvironment
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import java.io.InputStream
import java.nio.ByteBuffer
import java.util.*

/**
 * 检测结果数据类，封装物体检测后的输出图像和边界框信息
 * @property outputBitmap 处理后的输出图像（可能包含绘制的边界框）
 * @property outputBox 检测到的物体边界框数组，每个元素是包含边界框信息的Float数组
 *                     （通常包含中心点坐标、宽高、置信度、类别ID等）
 */
internal data class Result(
    var outputBitmap: Bitmap,
    var outputBox: Array<FloatArray>
)

/**
 * 物体检测器类，负责使用ONNX Runtime执行模型推理，处理输入输出数据
 * 核心功能是将输入图像转换为模型可处理的格式，运行推理并解析结果
 */
internal class ObjectDetector {

    /**
     * 执行物体检测的主要方法
     * @param inputStream 输入图像的输入流（通常来自assets或文件）
     * @param ortEnv ONNX Runtime环境对象，管理模型运行环境
     * @param ortSession ONNX Runtime会话对象，已加载模型并准备好推理
     * @return 包含检测结果的Result对象
     */
    fun detect(inputStream: InputStream, ortEnv: OrtEnvironment, ortSession: OrtSession): Result {
        // 步骤1：将输入图像的输入流转换为字节数组（原始图像字节数据）
        // 注意：这里假设输入流是原始图像数据（如JPEG/PNG格式），模型会自行解码
        val rawImageBytes = inputStream.readBytes()

        // 步骤2：定义输入张量的形状并创建ONNX张量
        // 形状为[图像字节数]，因为输入是原始图像字节流
        val shape = longArrayOf(rawImageBytes.size.toLong())

        // 创建输入张量，将字节数组包装为ByteBuffer，指定数据类型为UINT8（无符号8位整数）
        val inputTensor = OnnxTensor.createTensor(
            ortEnv,
            ByteBuffer.wrap(rawImageBytes),  // 包装原始图像字节
            shape,                           // 张量形状
            OnnxJavaType.UINT8               // 数据类型（图像字节通常为UINT8）
        )

        // 使用use块自动释放输入张量资源（避免内存泄漏）
        inputTensor.use {
            // 步骤3：调用ONNX Runtime会话执行模型推理
            // 输入映射：键为模型的输入名称（"image"，需与模型定义一致），值为输入张量
            // 输出集合：指定需要获取的输出名称（"image_out"和"scaled_box_out_next"，需与模型定义一致）
            val output = ortSession.run(
                Collections.singletonMap("image", inputTensor),
                setOf("image_out", "scaled_box_out_next")
            )

            // 使用use块自动释放输出资源
            output.use {
                // 步骤4：解析模型输出结果
                // 解析输出图像：第一个输出是处理后的图像字节数组
                val rawOutput = (output?.get(0)?.value) as ByteArray
                // 解析边界框：第二个输出是包含物体边界框信息的二维浮点数组
                val boxOutput = (output?.get(1)?.value) as Array<FloatArray>
                // 将图像字节数组转换为Bitmap对象
                val outputImageBitmap = byteArrayToBitmap(rawOutput)

                // 步骤5：封装并返回检测结果
                return Result(outputImageBitmap, boxOutput)
            }
        }
    }

    /**
     * 将字节数组转换为Bitmap图像
     * 用于将模型输出的图像字节数据转换为可在Android UI中显示的Bitmap
     * @param data 图像字节数组（通常为JPEG/PNG格式的原始字节）
     * @return 转换后的Bitmap对象
     */
    private fun byteArrayToBitmap(data: ByteArray): Bitmap {
        // 使用BitmapFactory解码字节数组为Bitmap
        // 参数：字节数组、起始偏移量（0表示从开头开始）、数据长度
        return BitmapFactory.decodeByteArray(data, 0, data.size)
    }
}